# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of MindSPONGE:
# MindSpore Simulation Package tOwards Next Generation molecular modelling.
#
# MindSPONGE is open-source software based on the AI-framework:
# MindSpore (https://www.mindspore.cn/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
Staging Tranformation PIMD integrator
"""

from typing import Tuple

import torch
import math
from torch import Tensor

from .integrator import Integrator, _integrator_register
from ..thermostat import Thermostat
from ..barostat import Barostat
from ..constraint import Constraint
from ...system import Molecule
from ...function import get_arguments, get_tensor


@_integrator_register('stgpimd')
class STGPIMD(Integrator):
    r"""A leap-frog integrator based on "middle scheme" developed by Jian Liu, et al. It is a subclass of `Integrator`.

    Reference:

        Zhang, Z.; Yan, K; Liu, X.; Liu, J.
        A Leap-Frog Algorithm-based Efficient Unified Thermostat Scheme for Molecular Dynamics [J].
        Chinese Science Bulletin, 2018, 63(33): 3467-3483.

    Args:

        system (Molecule):          Simulation system

        thermostat (Thermostat):    Thermostat for temperature coupling. Default: None

        barostat (Barostat):        Barostat for pressure coupling. Default: None

        constraint (Constraint):    Constraint algorithm. Default: None


    Supported Platforms:

        ``Ascend`` ``GPU``

    """
    def __init__(self,
                 system: Molecule,
                 temperature = 300.,
                 gamma0 = 10.,
                 remove_com = True,
                 method = 'BAOABnum',
                 **kwargs
                 ):

        super().__init__(
            system=system,
            thermostat=None,
            barostat=None,
            constraint=None,
        )
        self._kwargs = get_arguments(locals(), kwargs)
        self.name = 'PILE'
        n = self.num_walker
        self.ref_temp = self._get_mw_tensor(temperature, 'temperature')

        boltzmann = system.units.boltzmann
        hbar = 1.054572e-37 * 6.02e23 * 1e12
        self.beta = 1. / (boltzmann * self.ref_temp)
        self.omega_n = math.sqrt(n) / (self.beta * hbar) *torch.ones(n,1,1)

        self.gamma = math.sqrt(n) / (self.beta * hbar) *torch.ones(n,1,1)
        self.gamma0 = gamma0
        self.gamma[0] = self.gamma0
        self.c1 = torch.exp(-self.gamma*self.time_step)
        self.c2 = torch.sqrt(1-self.c1**2) / self.kinetic_unit_scale

        self.m_bar = torch.zeros_like(self.atom_mass.unsqueeze(-1)).repeat(n,1,1)
        self.m_ratio = torch.cat([torch.ones(1,), torch.arange(2,n+1) / torch.arange(1,n)]).unsqueeze(-1).unsqueeze(-1)
        self.m_bar[1:] = self.m_ratio[1:] * self.atom_mass.unsqueeze(-1)
        self.m_tilde = self.m_ratio * self.atom_mass.unsqueeze(-1)

        self._j = get_tensor(torch.arange(1,n+1).reshape(-1,1,1), torch.float32)
        self._ratio = (self._j-1) / self._j
        self.sqrt_m_tilde = torch.sqrt(self.m_tilde)
        self.remove_com = remove_com

        self._Fmatrix = get_tensor(torch.zeros((n,n)), torch.float32)
        for i in range(n):
            for j in range(i + 1):
                if i == j:
                    self._Fmatrix[i, j] = 1
                else:
                    self._Fmatrix[i, j] = j / i
        self._Fmatrix[0] = get_tensor(torch.ones(n), torch.float32)
        self._Fmatrix /= n

        self.method = method
        if self.method == 'BAOAB':
            self.a = torch.cos(self.omega_n*0.5*self.time_step)
            self.b = torch.sin(self.omega_n*0.5*self.time_step)/(self.omega_n*self.m_tilde)
            self.c = -torch.sin(self.omega_n*0.5*self.time_step)*self.omega_n*self.m_tilde
            self.d = torch.cos(self.omega_n*0.5*self.time_step)
            self.run = self.BAOAB
        elif self.method == 'BAOABnum':
            self.run = self.BAOABnum
        else:
            raise ValueError(f"Unsupported method: {self.method}")

    def print_info(self):
        info = {'Simulation Ensemble': 'NVT Ensemble',
                'Integrator Type': 'STG-PIMD',
                'time step': self.time_step,
                'temperature': self.ref_temp.item(),
                'gamma0': self.gamma0,}
        return info

    def forward(self,
               coordinate: Tensor,
               velocity: Tensor,
               force: Tensor,
               energy: Tensor,
               virial: Tensor = None,
               pbc_box: Tensor = None,
               steps: int = 0,
               **kwargs
                  ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]:

        variables = self.run(coordinate,
                            velocity,
                            force,
                            energy,
                            virial,
                            pbc_box,
                            steps,
                            **kwargs)
        return variables
 
    
    def langevin_thermostat(self, p, step):
        if self.control_step == 1 or step % self.control_step == 0:
            p = self.c1 * p + self.c2 * torch.sqrt(1 / self.beta) * self.sqrt_m_tilde * torch.randn(p.shape, device=p.device)
        return p
    
    def get_pimd_kinetics(self, p):
        ke = p ** 2 / (2*self.m_tilde)
        kinetics = torch.sum(ke, -2)
        return kinetics
    
    def get_pimd_com_velocity(self, momentum: Tensor):
        com_vel = momentum.sum(-2,keepdim=True)/self.m_tilde.sum(-2,keepdim=True) * self.m_tilde
        return com_vel
    
    def BAOABnum(self,
               coordinate: Tensor,
               velocity: Tensor,
               force: Tensor,
               energy: Tensor,
               virial: Tensor = None,
               pbc_box: Tensor = None,
               steps: int = 0,
               **kwargs
                  ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]:
        # (B,A,D) = (B,A,D) * (B,A,1)
        x_tilde = coordinate - self._ratio * coordinate.roll(-1,dims=0) - coordinate[[0]] /  self._j
        x_tilde[0] = coordinate[0]
        f_tilde = torch.einsum('jk,kad->jad',self._Fmatrix,force)

        p_tilde = velocity * self.m_tilde / torch.sqrt(self.m_ratio)
        p_tilde = p_tilde + self.acc_unit_scale * f_tilde * self.time_step - self.omega_n ** 2 * x_tilde * self.time_step * self.m_bar
        # (B,A,D) = (B,A,D) - (B,1,D)
        if self.remove_com:
            p_tilde = p_tilde - self.get_pimd_com_velocity(p_tilde)

        # R(t+0.5) = R(t) + v(t+0.5) * dt
        x_tilde = x_tilde + p_tilde / self.m_tilde * 0.5 * self.time_step

        if self.ref_temp is not None:
            # v'(t+0.5) = f_T[v(t+0.5)]
            p_tilde = self.langevin_thermostat(p_tilde, steps)

        # R(t+1) = R(t+0.5) + v'(t+0.5) * dt
        x_tilde = x_tilde + p_tilde / self.m_tilde * 0.5 * self.time_step
        velocity = p_tilde / self.m_tilde * torch.sqrt(self.m_ratio)

        summation_x = x_tilde / (self._j-1)
        summation_x[0] = 0
        coordinate = x_tilde[[0]] + (self._j-1) * torch.flip(torch.cumsum(torch.flip(summation_x,dims=[0]),0),dims=[0])

        variables = {'coordinate': coordinate,
                    'velocity': velocity,
                    'force': force,
                    'energy': energy,
                    'virial': virial,
                    'pbc_box': pbc_box,
                    }
        return variables
    
    def BAOAB(self,
               coordinate: Tensor,
               velocity: Tensor,
               force: Tensor,
               energy: Tensor,
               virial: Tensor = None,
               pbc_box: Tensor = None,
               steps: int = 0,
               **kwargs
                  ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]:
        
        # raise NotImplementedError
    
        x_tilde = coordinate - self._ratio * coordinate.roll(-1,dims=0) - coordinate[[0]] /  self._j
        x_tilde[0] = coordinate[0]
        f_tilde = torch.einsum('jk,kad->jad',self._Fmatrix,force)

        p_tilde = velocity * self.m_tilde / torch.sqrt(self.m_ratio)
        p_tilde = p_tilde + self.acc_unit_scale * f_tilde * self.time_step

        if self.remove_com:
            p_tilde = p_tilde - self.get_pimd_com_velocity(p_tilde)
        
        x_tilde_half = x_tilde.clone()
        p_tilde_half = p_tilde.clone()
        # R(t+0.5) = R(t) + v(t+0.5) * dt
        x_tilde_half[0] = x_tilde[0] + p_tilde[0] / self.m_tilde[0] * 0.5 * self.time_step
        x_tilde_half[1:] = self.a[1:] * x_tilde[1:] + self.b[1:] * p_tilde[1:]
        p_tilde_half[1:] = self.c[1:] * x_tilde[1:] + self.d[1:] * p_tilde[1:]

        if self.ref_temp is not None:
            # v'(t+0.5) = f_T[v(t+0.5)]
            p_tilde_half = self.langevin_thermostat(p_tilde_half, steps)
        
        x_tilde_new = x_tilde_half.clone()
        p_tilde_new = p_tilde_half.clone()
        # R(t+1) = R(t+0.5) + v'(t+0.5) * dt
        x_tilde_new[0] = x_tilde_half[0] + p_tilde_half[0] / self.m_tilde[0] * 0.5 * self.time_step
        x_tilde_new[1:] = self.a[1:] * x_tilde_half[1:] + self.b[1:] * p_tilde_half[1:]
        p_tilde_new[1:] = self.c[1:] * x_tilde_half[1:] + self.d[1:] * p_tilde_half[1:]

        summation_x = x_tilde_new / (self._j-1)
        summation_x[0] = 0
        coordinate = x_tilde_new[[0]] + (self._j-1) * torch.flip(torch.cumsum(torch.flip(summation_x,dims=[0]),0),dims=[0])
        velocity = p_tilde_new / self.m_tilde * torch.sqrt(self.m_ratio)

        variables = {'coordinate': coordinate,
                    'velocity': velocity,
                    'force': force,
                    'energy': energy,
                    'virial': virial,
                    'pbc_box': pbc_box,
                    }
        return variables

        



